import zarr
import numpy as np
import argparse
import os
import pickle
import json
from pathlib import Path
from vis_pcd import PointArrow, catpcd
from scalePcd import scalePcdD
from save_file import save_info
from scipy.spatial.transform import Rotation
from testTransform import transform_arrows
import open3d as o3d
from termcolor import cprint,colored
from num_print import p
def qxyzw_to_rpy(state_7d):
    assert len(state_7d) == 7, f"state_7d length: {len(state_7d)}"
    xyz = state_7d[:3]
    quat = state_7d[3:7]
    rpy = Rotation.from_quat([quat[1], quat[2], quat[3], quat[0]]).as_euler('xyz')
    return np.concatenate([xyz, rpy])
def rpy_to_qxyzw(state_6d):
    assert len(state_6d) == 6, f"state_6d length: {len(state_6d)}"
    xyz = state_6d[:3]
    rpy = state_6d[3:6]
    quat = Rotation.from_euler('xyz', rpy).as_quat()
    return np.concatenate([xyz, [quat[3], quat[0], quat[1], quat[2]]])
try:
    from franka_home_real import ARM_HOME
except ImportError:
    from .franka_home_real import ARM_HOME
class xyzOffset:
    def __init__(self):
        self.xyz_offset = np.zeros(3)
    def calibrate(self, xyz_zarr, xyz_real):
        self.xyz_offset = xyz_real - xyz_zarr
    def zarr_to_real(self, xyz_zarr):
        return xyz_zarr + self.xyz_offset
    def real_to_zarr(self, xyz_real):
        return xyz_real - self.xyz_offset
    def check_rpy(self, quat_zarr, rpy_real):
        rpy_zarr = Rotation.from_quat([quat_zarr[1], quat_zarr[2], quat_zarr[3], quat_zarr[0]]).as_euler('xyz')
        rpy_diff = rpy_real - rpy_zarr
        p.print("RPY difference", rpy_diff)
        return rpy_diff
# xyz_offset_instance = xyzOffset()
# xyz_offset_instance.calibrate(xyz_zarr, xyz_real)
# xyz_offset_instance.check_rpy(quat_zarr, rpy_real)
# xyz_quat[:3] = xyz_offset_instance.zarr_to_real(xyz_quat[:3])

class TrajOffsetParams:
    def __init__(self, config_path="/home/dwl/DemonGen_dev/real_world/traj.json"):
        with open(config_path, 'r') as f:
            config = json.load(f)
        pos = np.array([config["translation"]["x"], config["translation"]["y"], config["translation"]["z"]])
        self.TRAJ2POINTS_POS = pos
        self.pos_initial = self.TRAJ2POINTS_POS.copy()
        rpy_initial_deg = np.array([config["rotation"]["roll"], config["rotation"]["pitch"], config["rotation"]["yaw"]])
        self.rpy_initial = np.radians(rpy_initial_deg)
        self.offset_rpy = np.array([0.0, 0.0, 0.0])
        self.offset_pos = np.array([0.0, 0.0, 0.0])
        self.update_orientation()
        self.update_position()
    
    def update_orientation(self):
        self.TRAJ2POINTS_RPY = self.rpy_initial + self.offset_rpy
    
    def update_position(self):
        self.TRAJ2POINTS_POS = self.pos_initial + self.offset_pos
    
    def add_offset_rpy(self, delta_rpy, print_info=True):
        self.offset_rpy += delta_rpy
        self.update_orientation()
        if print_info:
            current_offset = self.offset_rpy.copy()
            current_offset_deg = np.degrees(current_offset)
            delta_rpy_deg = np.degrees(delta_rpy)
            print(colored(f"RPY offset applied. {delta_rpy_deg} degrees", "green"), colored(f"Current offset: {current_offset_deg} degrees", "green"))
        return self.offset_rpy
    
    def add_offset_pos(self, delta_pos, print_info=True):
        self.offset_pos += delta_pos
        self.update_position()
        if print_info:
            current_offset = self.offset_pos.copy()
            print(colored(f"Position offset applied.{delta_pos}", "green"), colored(f"Current offset: {current_offset}", "green"))
        return self.offset_pos
    
    def reset_offset(self):
        self.offset_rpy = np.array([0.0, 0.0, 0.0])
        self.offset_pos = np.array([0.0, 0.0, 0.0])
        self.update_orientation()
        self.update_position()
        return self.offset_rpy, self.offset_pos
    
    def get_offset_rpy(self):
        return self.offset_rpy.copy()
    
    def get_offset_pos(self):
        return self.offset_pos.copy()
    
    def get_pos(self):
        return self.TRAJ2POINTS_POS
    
    def get_rpy(self):
        return self.TRAJ2POINTS_RPY
    
    def save_calibration(self, config_path="/home/dwl/DemonGen_dev/real_world/traj.json"):
        final_pos = self.pos_initial + self.offset_pos
        final_rpy_rad = self.rpy_initial + self.offset_rpy
        final_rpy_deg = np.degrees(final_rpy_rad)
        calibration_data = {"translation": {"x": float(final_pos[0]), "y": float(final_pos[1]), "z": float(final_pos[2])}, "rotation": {"roll": float(final_rpy_deg[0]), "pitch": float(final_rpy_deg[1]), "yaw": float(final_rpy_deg[2])}}
        with open(config_path, 'w') as f:
            json.dump(calibration_data, f, indent=2)
        print(f"Calibration saved to {config_path}")
        print(f"Final translation: {final_pos}")
        print(f"Final rotation (RPY): {final_rpy_deg} degrees")
    
    def get_calibration_summary(self):
        final_pos = self.pos_initial + self.offset_pos
        final_rpy_rad = self.rpy_initial + self.offset_rpy
        final_rpy_deg = np.degrees(final_rpy_rad)
        initial_rpy_deg = np.degrees(self.rpy_initial)
        offset_rpy_deg = np.degrees(self.offset_rpy)
        summary = {"initial_translation": self.pos_initial.tolist(), "initial_rotation": initial_rpy_deg.tolist(), "current_offsets": {"position": self.offset_pos.tolist(), "rotation_rpy": offset_rpy_deg.tolist()}, "final_calibration": {"translation": final_pos.tolist(), "rotation_rpy": final_rpy_deg.tolist()}}
        cprint("Calibration Summary:", "blue")
        cprint(f"  Initial translation: {summary['initial_translation']}", "blue")
        cprint(f"  Initial rotation: {initial_rpy_deg} degrees", "blue")
        cprint(f"  Final translation: {summary['final_calibration']['translation']}", "green")
        cprint(f"  Final rotation: {final_rpy_deg} degrees", "green")
        initial_pos = np.array(summary['initial_translation'])
        final_pos = np.array(summary['final_calibration']['translation'])
        initial_rot = np.array(summary['initial_rotation'])
        final_rot = np.array(summary['final_calibration']['rotation_rpy'])
        pos_diff = final_pos - initial_pos
        rot_diff = final_rot - initial_rot
        cprint(f"  Position diff: {pos_diff}", "green")
        cprint(f"  Rotation diff: {rot_diff} degrees", "green")
        return summary

# Global instance
traj_offset_params = TrajOffsetParams()

class Traj2Points:
    def __init__(self, zarr_path):
        self.zarr_path = zarr_path
        self.zarr_array = zarr.open(zarr_path, mode='r')
        self.episode_ends = self.zarr_array['meta']['episode_ends'][:].tolist()
        self.states7D = self.zarr_array['ee']['state'][:,:7]#not used
        p.print("states7D shape", self.states7D.shape)
        # self.states7D = self.create_mimic_arrows_7d()
        self.states6D = np.array([qxyzw_to_rpy(state) for state in self.states7D])
        self.hand_states=self.zarr_array['ee']['state'][:,7]#not used
        first_timestep_state = self.states6D[0]
        self.xyz_zarr = first_timestep_state[:3]
        self.xyz_real = ARM_HOME[:3]
        self.rpy_zarr = first_timestep_state[3:6]
        self.rpy_real = ARM_HOME[3:6]
        self.joint_states8D = self.zarr_array['qpos']['state'][:,:8]
    def create_mimic_arrows_7d(self):
        arrow_direction = np.array([0, 0, 1])
        line_direction = np.array([1, 0, 0])
        arrow_positions = np.linspace([0.1, 0.1, 0.1], [0.3, 0.1, 0.1], 10)
        rotation = Rotation.align_vectors([[0, 0, 1]], [arrow_direction])[0]
        quat = rotation.as_quat()
        quat_xyzw = [quat[3], quat[0], quat[1], quat[2]]
        mimic_arrows = []
        for pos in arrow_positions:
            arrow_state = np.concatenate([pos, quat_xyzw])
            mimic_arrows.append(arrow_state)
        return np.array(mimic_arrows)
    def get_episode_timesteps(self, episode_idx):
        if isinstance(episode_idx, str):
            assert episode_idx.isdigit(), f"episode_idx: {episode_idx} is not a number"
            episode_idx = int(episode_idx)
        assert episode_idx >= 0, f"episode_idx: {episode_idx}"
        start_idx = 0 if episode_idx == 0 else self.episode_ends[episode_idx - 1]
        end_idx = self.episode_ends[episode_idx]
        return start_idx, end_idx
    def get_zarr_states(self, episode_idx,joint_state=False):
        start_idx, end_idx = self.get_episode_timesteps(episode_idx)
        if joint_state:
            return self.joint_states8D[start_idx:end_idx]
        else:
            return self.states6D[start_idx:end_idx]
    def get_zarr_hand_states(self, episode_idx):
        start_idx, end_idx = self.get_episode_timesteps(episode_idx)
        return self.hand_states[start_idx:end_idx]
    def state6d_to_arrows(self, states, start_color=[0, 0, 1], end_color=[1, 0, 0]):
        pointcloud = np.zeros((0, 6))
        num_timesteps = len(states)
        for timestep_idx, timestep_data in enumerate(states):
            if timestep_idx == 0:
                rpy_zarr = timestep_data[3:6]
                rpy_real = ARM_HOME[3:6]
            xyz_rpy = timestep_data[:6]
            time_ratio = timestep_idx / (num_timesteps - 1) if num_timesteps > 1 else 0
            color = [start_color[color_idx] * (1 - time_ratio) + end_color[color_idx] * time_ratio for color_idx in range(3)]
            arrow_points = PointArrow(xyz_rpy, color, 50,show_xy_arrow=True)
            pointcloud = catpcd(pointcloud, arrow_points)
        return pointcloud
     
    def get_pkl_pcd(self, pkl_path, timestep):
        assert os.path.exists(pkl_path), f"pkl_path: {pkl_path}"
        assert timestep >= 0, f"timestep: {timestep}"
        with open(pkl_path, 'rb') as f:
            demo_data = pickle.load(f)
        assert timestep < len(demo_data['point_cloud']), f"timestep: {timestep}, available: {len(demo_data['point_cloud'])}"
        frame_pcd = demo_data['point_cloud'][timestep]
        return frame_pcd
    def pcdOffset(self, pcd_frame):
        assert pcd_frame.shape[1] == 6, f"pcd_frame shape: {pcd_frame.shape}"
        xyz_offset = traj_offset_params.get_offset_pos() + self.get_ends_meet_xyz()
        rpy_offset_rad = traj_offset_params.get_offset_rpy() + self.get_ends_meet_rpy()
        rotation_matrix = Rotation.from_euler('xyz', rpy_offset_rad, degrees=False).as_matrix()
        xyz_coords = pcd_frame[:, :3] + xyz_offset
        rgb_values = pcd_frame[:, 3:6]
        center = self.xyz_zarr
        centered_xyz = xyz_coords - center
        rotated_centered_xyz = np.dot(centered_xyz, rotation_matrix.T)
        transformed_xyz = rotated_centered_xyz + center
        transformed_pcd = np.column_stack([transformed_xyz, rgb_values])
        return transformed_pcd
    def get_ends_meet_xyz(self):
        return self.xyz_zarr-self.xyz_real
    def get_ends_meet_rpy(self):
        p.print("rpy_zarr", self.rpy_zarr)
        p.print("rpy_real", self.rpy_real)
        return self.rpy_zarr-self.rpy_real
    def trajOffset(self, zarr_states):
        zarr_states = zarr_states.copy()
        assert len(zarr_states) > 0, f"zarr_states length: {len(zarr_states)}"
        zarr_states[:,0:3]=[zarr_state[:3]-self.get_ends_meet_xyz() for zarr_state in zarr_states]
        transformed_states = self.transform_arrows_array(zarr_states, self.xyz_real, traj_offset_params.get_offset_rpy()+self.get_ends_meet_rpy())
        p.print("zarr_states", zarr_states[0])
        p.print("transformed_states", transformed_states[0])
        transformed_states[:,0:3] += traj_offset_params.get_offset_pos()
        return np.array(transformed_states)

    def transform_arrows_array(self,arrows_array, center_point, offset_rpy_rad):
        p.print("arrows_array", arrows_array[1])
        p.print("center_point", center_point)
        p.print("offset_rpy_rad", offset_rpy_rad)
        arrows_array=np.array(arrows_array)
        center_point=np.array(center_point)
        offset_rpy_rad=-np.array(offset_rpy_rad)
        assert arrows_array.shape[1] == 6, f"arrows_array shape error: {arrows_array.shape}"
        assert center_point.shape == (3,), f"center_point shape error: {center_point.shape}"
        assert offset_rpy_rad.shape == (3,), f"offset_rpy_rad shape error: {offset_rpy_rad.shape}"
        roll_rad, pitch_rad, yaw_rad = offset_rpy_rad
        center_x, center_y, center_z = center_point
        rotation_matrix = np.array([
            [np.cos(yaw_rad) * np.cos(pitch_rad), np.cos(yaw_rad) * np.sin(pitch_rad) * np.sin(roll_rad) - np.sin(yaw_rad) * np.cos(roll_rad), np.cos(yaw_rad) * np.sin(pitch_rad) * np.cos(roll_rad) + np.sin(yaw_rad) * np.sin(roll_rad)],
            [np.sin(yaw_rad) * np.cos(pitch_rad), np.sin(yaw_rad) * np.sin(pitch_rad) * np.sin(roll_rad) + np.cos(yaw_rad) * np.cos(roll_rad), np.sin(yaw_rad) * np.sin(pitch_rad) * np.cos(roll_rad) - np.cos(yaw_rad) * np.sin(roll_rad)],
            [-np.sin(pitch_rad), np.cos(pitch_rad) * np.sin(roll_rad), np.cos(pitch_rad) * np.cos(roll_rad)]
        ])
        transformed_arrows = np.zeros_like(arrows_array)
        for arrow_index in range(arrows_array.shape[0]):
            arrow_xyz = arrows_array[arrow_index, :3] - center_point
            arrow_rpy = arrows_array[arrow_index, 3:]
            transformed_xyz = np.dot(rotation_matrix, arrow_xyz) + center_point
            transformed_rpy = arrow_rpy + offset_rpy_rad
            transformed_arrows[arrow_index, :3] = transformed_xyz
            transformed_arrows[arrow_index, 3:] = transformed_rpy
        return transformed_arrows
    def test_pcdOffset(self, episode_idx, pkl_path, timestep, print_color="black"):
        zarr_states = self.get_zarr_states(episode_idx)
        p.print("pcdOffset zarr_states", zarr_states[0][3:6])
        zarr_pcd = self.state6d_to_arrows(zarr_states)
        pkl_pcd = self.get_pkl_pcd(pkl_path, timestep)
        transformed_pkl_pcd = self.pcdOffset(pkl_pcd)
        offset_pos = traj_offset_params.get_offset_pos()
        offset_rpy = traj_offset_params.get_offset_rpy()
        combined_pcd = catpcd(zarr_pcd, transformed_pkl_pcd)
        save_info("pcd", f"{os.path.basename(self.zarr_path)}/episode{episode_idx}_{'_'.join([str(round(ri,2)) for ri in np.concatenate([offset_pos, offset_rpy])])}", combined_pcd,print_color=print_color)
        return combined_pcd
    def test_trajOffset(self, episode_idx, pkl_path, timestep, print_color="black"):
        zarr_states = self.get_zarr_states(episode_idx)
        transformed_zarr_states = self.trajOffset(zarr_states)
        p.print("traj", transformed_zarr_states[0][3:6])
        transformed_zarr_pcd = self.state6d_to_arrows(transformed_zarr_states)
        pkl_pcd = self.get_pkl_pcd(pkl_path, timestep)
        combined_pcd = catpcd(transformed_zarr_pcd, pkl_pcd)
        offset_pos = traj_offset_params.get_offset_pos()
        offset_rpy = traj_offset_params.get_offset_rpy()
        save_info("pcd", f"{os.path.basename(self.zarr_path)}/episode{0}_traj_{'_'.join([str(round(ri,2)) for ri in np.concatenate([offset_pos, offset_rpy])])}", combined_pcd,print_color=print_color)
        return combined_pcd
    def get_qpos(self,episode_idx):
        zarr_states = self.get_zarr_states(episode_idx,joint_state=True)
        return zarr_states
    def get_zarr_local(self,episode_idx):
        zarr_states = self.get_zarr_states(episode_idx)
        zarr_hand_states = self.get_zarr_hand_states(episode_idx).reshape(-1,1)
        zarr_local_states = self.trajOffset(zarr_states)
        p.print("zarr_local_states", zarr_local_states[0])
        transformed_zarr_pcd = self.state6d_to_arrows(zarr_local_states)
        return np.concatenate([zarr_local_states, zarr_hand_states], axis=1),transformed_zarr_pcd
    
if __name__ == "__main__":
    # Test TrajOffsetParams
    # print("Testing TrajOffsetParams:")
    # traj_offset_params.add_offset_pos([0.01, 0.02, 0.03])
    # traj_offset_params.add_offset_rpy([0.1, 0.2, 0.3])
    # traj_offset_params.get_calibration_summary()
    
    zarr_path = f"{Path(__file__).parent.parent}/zarrtraj/grab-roller-demotest.zarr"
    p.print("zarr_path", zarr_path)
    traj_processor = Traj2Points(zarr_path)
    pkl_path = f"{Path(__file__).parent.parent}/data/source_demos/name/1.pkl"
    
    # Test different yaw offsets using trajOffsetParams
    for offset_rpy in [[0,0,0],[0,0,45],[0,0,90],[0,0,135],[0,0,180]]: 
        traj_offset_params.reset_offset()
        traj_offset_params.add_offset_rpy(np.radians(offset_rpy))
        # traj_processor.test_pcdOffset(0, pkl_path, 0, print_color="red")
        traj_processor.test_trajOffset(0, pkl_path, 0, print_color="blue")

